// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/sparksql/aggregates/register.h>

using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::functions::aggregate::sparksql::test {

namespace {
class CentralMomentsAggregationTest : public AggregationTestBase {
 protected:
  void SetUp() override {
    AggregationTestBase::SetUp();
    registerAggregateFunctions("spark_");
  }

  void testCenteralMomentsAggResult(
      const std::string& agg,
      const RowVectorPtr& input,
      const RowVectorPtr& expected,
      bool legacy = false) {
    auto plan = PlanBuilder()
                    .values({input})
                    .singleAggregation({}, {fmt::format("spark_{}(c0)", agg)})
                    .planNode();
    AssertQueryBuilder(plan)
        .config(
            core::QueryConfig::kSparkLegacyStatisticalAggregate,
            legacy ? "true" : "false")
        .assertResults({expected});
  }
};

TEST_F(CentralMomentsAggregationTest, skewnessHasResult) {
  auto agg = "skewness";
  auto input = make_row_vector({make_flat_vector<double>({1, 2})});
  // Even when the count is 2, Spark still produces output.
  auto expected =
      make_row_vector({make_flat_vector<double>(std::vector<double>{0.0})});
  testCenteralMomentsAggResult(agg, input, expected);

  input = make_row_vector({make_flat_vector<double>({1, 1})});
  expected = make_row_vector({make_nullable_flat_vector<double>(
      std::vector<std::optional<double>>{std::nullopt})});
  testCenteralMomentsAggResult(agg, input, expected);

  // Output NULL when m2 equals 0.
  input = make_row_vector({make_flat_vector<double>({1, 1})});
  expected = make_row_vector({make_nullable_flat_vector<double>(
      std::vector<std::optional<double>>{std::nullopt})});
  testCenteralMomentsAggResult(agg, input, expected);

  // Output NaN when m2 equals 0 for legacy aggregate.
  input = make_row_vector({make_flat_vector<double>({1, 1})});
  expected = make_row_vector({make_flat_vector<double>(
      std::vector<double>{std::numeric_limits<double>::quiet_NaN()})});
  testCenteralMomentsAggResult(agg, input, expected, true);
}

TEST_F(CentralMomentsAggregationTest, pearsonKurtosis) {
  auto agg = "kurtosis";
  auto input = make_row_vector({make_flat_vector<double>({1, 10, 100, 10, 1})});
  auto expected = make_row_vector(
      {make_flat_vector<double>(std::vector<double>{0.19432323191699075})});
  testCenteralMomentsAggResult(agg, input, expected);

  input = make_row_vector({make_flat_vector<double>({-10, -20, 100, 1000})});
  expected = make_row_vector(
      {make_flat_vector<double>(std::vector<double>{-0.7014368047529627})});
  testCenteralMomentsAggResult(agg, input, expected);

  // Even when the count is 2, Spark still produces non-null result.
  input = make_row_vector({make_flat_vector<double>({1, 2})});
  expected = make_row_vector({make_flat_vector<double>(std::vector<double>{-2.0})});
  testCenteralMomentsAggResult(agg, input, expected);

  // Output NULL when m2 equals 0.
  input = make_row_vector({make_flat_vector<double>({1, 1})});
  expected = make_row_vector({make_nullable_flat_vector<double>(
      std::vector<std::optional<double>>{std::nullopt})});
  testCenteralMomentsAggResult(agg, input, expected);

  // Output NaN when m2 equals 0 for legacy aggregate.
  input = make_row_vector({make_flat_vector<double>({1, 1})});
  expected = make_row_vector({make_flat_vector<double>(
      std::vector<double>{std::numeric_limits<double>::quiet_NaN()})});
  testCenteralMomentsAggResult(agg, input, expected, true);
}

} // namespace
} // namespace kumo::pollux::functions::aggregate::sparksql::test
